package cn.ondraw.ragflow.service.impl;


import cn.hutool.json.JSONUtil;
import cn.ondraw.ragflow.pojo.request.*;
import cn.ondraw.ragflow.pojo.response.*;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import cn.ondraw.ragflow.pojo.request.*;
import cn.ondraw.ragflow.pojo.response.*;
import cn.ondraw.ragflow.pojo.request.*;
import cn.ondraw.ragflow.pojo.response.*;
import cn.ondraw.ragflow.service.ChunkService;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Service;

import java.util.Optional;
import java.util.StringJoiner;

@Service
@Slf4j
@RequiredArgsConstructor
public class ChunkServiceImpl implements ChunkService {
    private final PythonApiCaller pythonApiCaller;

    private static final String BASE_URL = "/api/v1/datasets";

    /**
     * 添加块到指定数据集和文档中
     *
     * @param datasetId  数据集 ID
     * @param documentId 文档 ID
     * @param request    添加块的请求对象
     * @return 添加块的响应对象
     */
    @Override
    public AddChunkResponse addChunk(String datasetId, String documentId, AddChunkRequest request) {
        String url = String.format("%s/%s/documents/%s/chunks", BASE_URL, datasetId, documentId);
        ResponseEntity<?> response = pythonApiCaller.callPostApi(url, request);
        return JSONUtil.toBean(response.getBody() != null ? response.getBody().toString() : null, AddChunkResponse.class, true);
    }

    /**
     * 列出指定数据集和文档中的块
     *
     * @param request 列出块的请求对象
     * @return 列出块的响应对象
     */
    @Override
    public ListChunksResponse listChunks(ListChunksRequest request) {
        StringJoiner params = new StringJoiner("&");
        Optional.ofNullable(request.getKeywords()).ifPresent(v -> params.add("keywords=" + v));
        Optional.ofNullable(request.getPage()).ifPresent(v -> params.add("page=" + v));
        Optional.ofNullable(request.getPage_size()).ifPresent(v -> params.add("page_size=" + v));
        Optional.ofNullable(request.getId()).ifPresent(v -> params.add("id=" + v));

        String url = String.format("/datasets/%s/documents/%s/chunks%s", request.getDataset_id(), request.getDocument_id(), params.length() > 0 ? "?" + params : "");
        ResponseEntity<?> response = pythonApiCaller.callGetApi(url);
        return JSONUtil.toBean(response.getBody() != null ? response.getBody().toString() : null, ListChunksResponse.class, true);
    }

    /**
     * 删除指定数据集和文档中的块
     *
     * @param datasetId  数据集 ID
     * @param documentId 文档 ID
     * @param request    删除块的请求对象
     * @return 删除块的响应对象
     */
    @Override
    public DeleteChunksResponse deleteChunks(String datasetId, String documentId, DeleteChunksRequest request) {
        String url = String.format("/%s/%s/documents/%s/chunks", BASE_URL, datasetId, documentId);
        ResponseEntity<?> response = pythonApiCaller.callDeleteApi(url, request);
        return JSONUtil.toBean(response.getBody() != null ? response.getBody().toString() : null, DeleteChunksResponse.class, true);
    }

    /**
     * 更新指定数据集、文档和块 ID 的块信息
     *
     * @param datasetId  数据集 ID
     * @param documentId 文档 ID
     * @param chunkId    块 ID
     * @param request    更新块的请求对象
     * @return 更新块的响应对象
     */
    @Override
    public UpdateChunkResponse updateChunk(String datasetId, String documentId, String chunkId, UpdateChunkRequest request) {
        String url = String.format("/datasets/%s/documents/%s/chunks/%s", datasetId, documentId, chunkId);
        ResponseEntity<?> response = pythonApiCaller.callPutApi(url, request);
        return JSONUtil.toBean(response.getBody() != null ? response.getBody().toString() : null, UpdateChunkResponse.class, true);
    }

    /**
     * 检索符合条件的块
     *
     * @param request 检索块的请求对象
     * @return 检索块的响应对象
     */
    @Override
    public RetrieveChunksResponse retrieveChunks(RetrieveChunksRequest request) {
        String url = "/retrieval";
        ResponseEntity<?> response = pythonApiCaller.callPostApi(url, request);
        return JSONUtil.toBean(response.getBody() != null ? response.getBody().toString() : null, RetrieveChunksResponse.class, true);
    }

    /**
     * 测试检索块
     *
     * @param request 检索块的请求对象
     * @return 检索块的响应对象
     */
    @Override
    public Object chunksRetrieveTest(RetrieveTestChunksRequest request) {
        String url = "/v1/chunk/retrieval_test";
        ResponseEntity<?> response = pythonApiCaller.callPostApi(url, request, false);
        return JSONUtil.toBean(response.getBody() != null ? response.getBody().toString() : null, Object.class, true);
    }
}
